import json
from config import Config
from bert4vec import Bert4Vec

model = Bert4Vec(mode='roformer-sim-small')
with open(Config.kg, encoding="utf8") as f:
    data = json.load(f)
result = [element["question"] for i, element in enumerate(data)]
f = open(Config.sentences_path, "w", encoding="utf8")
for element in list(result):
    f.write(element + "\n")
f.close()
model.build_index(Config.sentences_path, ann_search=True, gpu_index=Config.ISGPU, n_search=32)


def get_response(text):
    """
    问答相似度匹配获取匹配的问题及答案
    """
    results_ = model.search(queries=[text], threshold=Config.threshold, top_k=5)
    if results_:
        results = results_[0][0][0]

        thresh = results_[0][0][1]
        answer = data[result.index(results)]["answer"]
        return results, answer, thresh

    else:
        return "", Config.UNKNOW, 1.0


if __name__ == '__main__':
    print(get_response("居家隔离要注意什么"))
